#include "cu_sftp.h"

#define SFTP_SERVER_URL_HEADER  "sftp://"

cu_sftp_t *cu_sftp_init(const char *ip, const char *port, const char *username, const char *password) {
    cu_sftp_t *cu_sftp_handle = (cu_sftp_t *)malloc(sizeof(cu_sftp_t));
    if (cu_sftp_handle == NULL) {
        printf("Failed to allocate memory for cu_sftp_handle.\n");
        return NULL;
    }

    cu_sftp_handle->ip = strdup(ip);
    cu_sftp_handle->port = strdup(port);
    cu_sftp_handle->username = strdup(username);
    cu_sftp_handle->password = strdup(password);

    return cu_sftp_handle;
}

static size_t sftp_read_data(void *ptr, size_t size, size_t nmemb, FILE *stream) {
    size_t read_bytes = fread(ptr, size, nmemb, stream);
    return read_bytes;
}

static size_t sftp_write_data(void *ptr, size_t size, size_t nmemb, FILE *stream) {
    size_t written_bytes = fwrite(ptr, size, nmemb, stream);
    return written_bytes;
}

static int sftp_progress_callback(void* clientp, double dltotal, double dlnow, double ultotal, double ulnow)
{
	if (dltotal > 0){
		int nPersent = (int)(100.0 * dlnow / dltotal);
		printf("[%f/%f] download progress:%d\r", dlnow, dltotal, nPersent);
	}
	return 0;
}

static curl_off_t sftp_get_remote_file_size(cu_sftp_t * cu_sftp_handle, const char *remote_file_path) {
    CURL *curl;
    CURLcode res;
    curl_off_t remote_file_size = -1;

    // 初始化libcurl
    curl_global_init(CURL_GLOBAL_DEFAULT);
    curl = curl_easy_init();
    if(!curl) {
        return remote_file_size;
    }
    curl_easy_setopt(curl, CURLOPT_VERBOSE, 1L);
    // 首先尝试获取远程文件的大小
    size_t url_len = strlen(SFTP_SERVER_URL_HEADER) + strlen(cu_sftp_handle->ip) + strlen(cu_sftp_handle->port) + strlen(remote_file_path) + 2;
    char *url = (char *)malloc(url_len);
    snprintf(url, url_len, "%s%s:%s%s", SFTP_SERVER_URL_HEADER, cu_sftp_handle->ip, cu_sftp_handle->port, remote_file_path);
    curl_easy_setopt(curl, CURLOPT_URL, url);
    curl_easy_setopt(curl, CURLOPT_USERNAME, cu_sftp_handle->username);
    curl_easy_setopt(curl, CURLOPT_PASSWORD, cu_sftp_handle->password);
    curl_easy_setopt(curl, CURLOPT_NOBODY, 1L); // HEAD request to get file size
    curl_easy_setopt(curl, CURLOPT_HEADERFUNCTION, NULL);
    curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, NULL);
    curl_easy_setopt(curl, CURLOPT_FILETIME, 1L);

    res = curl_easy_perform(curl);
    if(res == CURLE_OK) {
        // 获取远程文件大小
        res = curl_easy_getinfo(curl, CURLINFO_CONTENT_LENGTH_DOWNLOAD_T, &remote_file_size);
        if(res == CURLE_OK) {
            printf("remote file size: %ld\n", remote_file_size);
        }else {
            fprintf(stderr, "Failed to get remote file size: %s\n", curl_easy_strerror(res));
            remote_file_size = -1;
        }
    }else {
        fprintf(stderr, "Failed to get remote file size: %s, %d\n", curl_easy_strerror(res), res);
        if(res == CURLE_REMOTE_FILE_NOT_FOUND) {
            remote_file_size = 0; 
        }else {
            remote_file_size = -1;
        }
    }

    curl_easy_cleanup(curl);
    free(url);

    return remote_file_size;
}

bool cu_sftp_upload(cu_sftp_t * cu_sftp_handle, const char *remote_file_path, const char *local_file_path, bool resume) {
    CURL *curl;
    CURLcode res;
    FILE *local_file;
    long local_file_size = 0;
    long remote_file_size = 0;

    // 初始化libcurl
    curl_global_init(CURL_GLOBAL_DEFAULT);
    curl = curl_easy_init();
    if(!curl) {
        printf("Failed to init curl.\n");
        return false;
    }
    local_file = fopen(local_file_path, "rb");
    if(!local_file) {
        printf("Failed to open local file.\n");
        curl_global_cleanup();
        return false;
    }
    if(resume) {
        // 获取本地文件大小
        fseek(local_file, 0, SEEK_END);
        local_file_size = ftell(local_file);
        fseek(local_file, 0, SEEK_SET);
        // 获取远程文件大小
        remote_file_size = sftp_get_remote_file_size(cu_sftp_handle, remote_file_path);
        if(remote_file_size == -1) {
            fclose(local_file);
            curl_global_cleanup();
            return false;
        }
        // 如果远程文件存在且大小小于本地文件，设置断点续传
        if(remote_file_size > 0) {
            if(remote_file_size < local_file_size) {
                fseek(local_file, remote_file_size, SEEK_SET); // 跳过已上传部分
            }else {
                fprintf(stderr, "remote file size is larger or equal than local file size, no need to upload.\n");
                fclose(local_file);
                curl_global_cleanup();
                return true;
            }
        } else {
            fprintf(stderr, "transfer from head\n");
            remote_file_size = 0; // 从头开始上传
        }
    }else {
        remote_file_size = 0; // 从头开始上传
    }

    curl_easy_setopt(curl, CURLOPT_VERBOSE, 1L);
    // 首先尝试获取远程文件的大小
    size_t url_len = strlen(SFTP_SERVER_URL_HEADER) + strlen(cu_sftp_handle->ip) + strlen(cu_sftp_handle->port) + strlen(remote_file_path) + 2;
    char *url = (char *)malloc(url_len);
    snprintf(url, url_len, "%s%s:%s%s", SFTP_SERVER_URL_HEADER, cu_sftp_handle->ip, cu_sftp_handle->port, remote_file_path);
    // 设置上传选项
    curl_easy_setopt(curl, CURLOPT_URL, url);
    curl_easy_setopt(curl, CURLOPT_USERNAME, cu_sftp_handle->username);
    curl_easy_setopt(curl, CURLOPT_PASSWORD, cu_sftp_handle->password);
    curl_easy_setopt(curl, CURLOPT_UPLOAD, 1L);
    curl_easy_setopt(curl, CURLOPT_READFUNCTION, sftp_read_data);
    curl_easy_setopt(curl, CURLOPT_READDATA, local_file);
    curl_easy_setopt(curl, CURLOPT_INFILESIZE_LARGE, (curl_off_t)local_file_size);

    // 如果是断点续传，设置上传起始位置
    if(remote_file_size > 0) {
        curl_easy_setopt(curl, CURLOPT_APPEND, 1L);
    }

    // 执行上传
    res = curl_easy_perform(curl);
    // 检查上传结果
    if(res != CURLE_OK) {
        fprintf(stderr, "curl_easy_perform() failed: %s\n", curl_easy_strerror(res));
    } else {
        printf("File uploaded successfully.\n");
    }

    // 清理
    fclose(local_file);
    curl_easy_cleanup(curl);
    free(url);
    return res == CURLE_OK ? true : false;
}

bool cu_sftp_download(cu_sftp_t * cu_sftp_handle, const char *remote_file_path, const char *local_file_path) {
    CURL *curl;
    CURLcode res;
    FILE *local_file;

    curl_global_init(CURL_GLOBAL_DEFAULT);
    curl = curl_easy_init();
    if(!curl) {
        printf("Failed to init curl.\n");
        return false;
    }
    local_file = fopen(local_file_path, "wb");
    if(!local_file) {
        printf("Failed to open local file.\n");
        curl_global_cleanup();
        return false;
    }
    curl_easy_setopt(curl, CURLOPT_VERBOSE, 1L);
    // 首先尝试获取远程文件的大小
    size_t url_len = strlen(SFTP_SERVER_URL_HEADER) + strlen(cu_sftp_handle->ip) + strlen(cu_sftp_handle->port) + strlen(remote_file_path) + 2;
    char *url = (char *)malloc(url_len);
    snprintf(url, url_len, "%s%s:%s%s", SFTP_SERVER_URL_HEADER, cu_sftp_handle->ip, cu_sftp_handle->port, remote_file_path);
    printf("url: %s\n", url);
    // 设置上传选项
    curl_easy_setopt(curl, CURLOPT_URL, url);
    curl_easy_setopt(curl, CURLOPT_USERNAME, cu_sftp_handle->username);
    curl_easy_setopt(curl, CURLOPT_PASSWORD, cu_sftp_handle->password);
    curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, sftp_write_data);
    curl_easy_setopt(curl, CURLOPT_WRITEDATA, local_file);

    // 显示重定向次数
    curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);
    // 是否验证证书
    curl_easy_setopt(curl, CURLOPT_SSL_VERIFYPEER, 0L);
    // 是否验证主机名
    curl_easy_setopt(curl, CURLOPT_SSL_VERIFYHOST, 0L);

    curl_easy_setopt(curl, CURLOPT_FTP_SKIP_PASV_IP, 1L); // add
    curl_easy_setopt(curl, CURLOPT_TCP_KEEPALIVE, 1L);    // add
    curl_easy_setopt(curl, CURLOPT_SSL_VERIFYPEER, 0L);
    curl_easy_setopt(curl, CURLOPT_SSL_VERIFYHOST, 0L);
    curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);
    curl_easy_setopt(curl, CURLOPT_BUFFERSIZE, 102400L);
    curl_easy_setopt(curl, CURLOPT_MAXREDIRS, 50L);

        // 设置文件下载进度回调
    curl_easy_setopt(curl, CURLOPT_NOPROGRESS, 0L);
    // 设置进度回调函数
    curl_easy_setopt(curl, CURLOPT_PROGRESSFUNCTION, sftp_progress_callback);

    #ifndef DISABLE_SSH_AGENT
        curl_easy_setopt(curl, CURLOPT_SSH_AUTH_TYPES, CURLSSH_AUTH_AGENT);
    #endif

    res = curl_easy_perform(curl);
    // 检查上传结果
    if(res != CURLE_OK) {
        fprintf(stderr, "curl_easy_perform() failed: %s\n", curl_easy_strerror(res));
    } else {
        printf("File download successfully.\n");
    }
    fclose(local_file);
    curl_global_cleanup();
    free(url);
    
    return res == CURLE_OK ? true : false;
}

bool cu_sftp_destroy(cu_sftp_t * cu_sftp_handle) {
    if (cu_sftp_handle == NULL) {
        return false;
    }

    free(cu_sftp_handle->ip);
    free(cu_sftp_handle->port);
    free(cu_sftp_handle->username);
    free(cu_sftp_handle->password);
    free(cu_sftp_handle);

    return true;
}
